#!/usr/bin/env python3
# -*- coding:utf-8 -*-

'''TODO: use treenode?'''

import os
import sys
import argparse
import subprocess

from ctools.subcmd.ctbase import CToolsCmd

from ctools.gvhelper.graphwrapper import GraphWrapper
import ctools.gvhelper.render as gvrender

import logging
from ctools.logext import *

_fpath = os.path.basename(os.path.realpath(__file__))
_rlog, _plog, _mlog, _tlog = get_loggers(__file__, __name__, propagate=False)

def check_git_repo(path:str) -> bool:
    cmd = f'git -C {path} rev-parse --show-toplevel'
    ret = -1
    try:
        ret = subprocess.run(cmd, stdout=subprocess.PIPE, cwd=path, shell=True, text=True)
    finally:
        return True if ret == 0 else False

def do_grep(path, is_git_repo=False):
    cmd = r"""grep -rsIw '^[ \t]*SYSCALL_DEFINE[0-6]' arch/arm64 [b-z]*/"""
    if is_git_repo:
        cmd = r"""git grep -Iw '^[ \t]*SYSCALL_DEFINE[0-6]' arch/arm64 [b-z]*/"""

    try:
        ret = subprocess.run(cmd, stdout=subprocess.PIPE, cwd=path, shell=True, text=True)
        return ret.stdout
    except Exception as e:
        print(e)

_ignores = ['nommu.c', 'posix-stubs.c']

def process_lines(lines):
    subgs = dict()
    for i, line in enumerate(lines):
        line = line.strip()
        if not line:
            continue
        if line.count(':') < 1:
            _plog.warning(line)
            continue

        strs = line.split(':', 1)
        fpath = strs[0]
        pos1 = strs[1].find('(')
        pos = strs[1].find(',')
        if pos >= 0:
            scall = strs[1][pos1 + 1:pos]
        else:
            scall = strs[1][pos1 + 1:].strip('()')
        scall = scall.strip()

        scnr = strs[1].split('(')[0]
        scnr = scnr.replace('SYSCALL_DEFINE', '').strip()

        scall = (scall, scnr)

        folder = os.path.dirname(fpath)
        cfile = os.path.basename(fpath)
        if cfile in _ignores:
            continue
        if not cfile.endswith('.c'):
            continue

        tmp = subgs
        if '/' in fpath:
            for s in folder.split('/'):
                if s in tmp.keys():
                    tmp = tmp[s]
                else:
                    tmp[s] = dict()
                    tmp = tmp[s]

        if cfile in tmp.keys():
            tmp[cfile].append(scall)
        else:
            tmp[cfile] = [scall]
    return subgs

def _func(prefix, d):
    buf = ''
    cnt = 0
    for k in d.keys():
        tmp = d[k]
        if type(tmp) == dict:
            sg = "{}_{}".format(prefix, k)
            buf += sg + ' {\n'
            l = sg
            _tmp, _cnt = _func(sg, tmp)

            l = 'label="{} ({})"'.format(
                l.replace('subgraph cluster_', '').strip('_').replace('_', '/'),
                _cnt)
            l += ';\n'
            buf += l
            buf += _tmp
            cnt += _cnt
            buf += "}\n"
        else:
            sg = "{}_{}".format(prefix, k)
            buf += sg.replace('.', '_').replace('-', '_') + ' {\n'
            buf += 'label="{}({})"\n'.format(k, len(tmp))
            cnt += len(tmp)
            for sc in tmp:
                buf += '{}[label="{}({})"];\n'.format(sc[0], sc[0], sc[1])
            buf += "}\n"
    return buf, cnt

def gen_gvtext(subgs) -> str:
    _tmp, _cnt =  _func('subgraph cluster', subgs)

    buf = "digraph {\n"
    buf += "    layout=fdp\n"
    buf += "    rankdir=LR;\n"
    buf += "    node [shape=box];\n"
    buf += f'    label="ioctls ({_cnt})"\n'

    buf += _tmp
    buf += "}\n"
    return buf


class SyscallCmd(CToolsCmd):
    @classmethod
    def name(cls):
        return "syscall"

    @classmethod
    def description(cls):
        return "syscall graph tool"

    def do_command(self, argv:list[str]):
        argp = argparse.ArgumentParser(description='syscall view tool')
        argp.add_argument("-e", "--engine", type=str, default="dot")
        argp.add_argument("-o", "--output", type=str, required=True)
        argp.add_argument("-O", "--option", type=str, default='', choices=('o', 'r', 'or',))
        argp.add_argument("-f", "--format", type=str, default=None)
        argp.add_argument('kernel_path', type=str)

        args = argp.parse_args(argv)

        buf = do_grep(args.kernel_path, check_git_repo(args.kernel_path))
        text = gen_gvtext(process_lines(buf.splitlines()))

        if 'o' in args.option:
            with open(args.output + '_o.gv', "w") as f:
                f.write(text)

        graph = GraphWrapper(string=text)
        _render = gvrender.GvRender()
        _render.render_all(graph)

        if 'r' in args.option:
            graph.write(args.output + '_r.gv')

        graph.draw(filename=args.output, prog=args.engine, format=args.format)

if __name__ == '__main__':
    SyscallCmd().do_command(sys.argv[1:])

